import gzip
import os
import re
import vcf


def create_snv_vep_input(in_vcf: str, vep_input: str):
    with open(vep_input, 'w') as writer:
        writer.write('#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\n')
        reader = vcf.Reader(filename=in_vcf, encoding='utf-8')
        for row in reader:
            if min(len(row.REF), len(row.ALT[0])) == 1:
                key = f'{row.CHROM}-{row.POS}-{row.REF}-{row.ALT[0]}'
                writer.write(f'{row.CHROM}\t{row.POS}\t{key}\t{row.REF}\t{row.ALT[0]}\t.\tPASS\t.\n')


def create_cnv_vep_input(in_vcf: str, vep_input: str):
    with open(vep_input, 'w') as writer:
        reader = vcf.Reader(filename=in_vcf, encoding='utf-8')
        for row in reader:
            chrom, start, end, alt = str(row.CHROM), int(row.POS), re.sub('<|>', '', str(row.ALT[0])), row.INFO.get('END')
            writer.write(f'{chrom}\t{start}\t{end}\t{alt}')


def read_vep_out(vep_out: str) -> dict:
    var_dict = dict()
    with open(vep_out) as reader:
        header = list()
        for line in reader:
            line = line.strip()
            if line.startswith("#"):
                if not line.startswith("##"):
                    header = line.lstrip('#').split("\t")
                continue
            records = line.strip().split("\t")
            row = dict(zip(header, records))
            var_dict.setdefault(row['Uploaded_variation'], list()).append(row)
    return var_dict


def run_vep_snv(vep_input: str, vep_output: str,  vep_cache: str,  vep_assembly: str, process: int):
    vep_command = '''
            vep --offline --refseq --use_given_ref \
            --dir_cache ''' + vep_cache + ''' \
            --species "homo_sapiens" \
            --assembly ''' + vep_assembly + ''' \
            --fork  ''' + str(process) + ''' \
            --canonical \
            --flag_pick \
            --hgvs --hgvsg --symbol \
            --distance 500 \
            --exclude_predicted \
            --lookup_ref \
            --numbers \
            --force \
            --input_file ''' + vep_input + '''\
            --output_file ''' + vep_output + ''' --no_stats \
            --tab --fields "Uploaded_variation,SYMBOL,Feature,CANONICAL,PICK,Consequence,HGVSc,HGVSp,HGVSg,EXON,INTRON"
        '''
    if os.system(vep_command) > 0:
        raise Exception('run VEP FAILED')


def run_vep_cnv(vep_input: str, vep_output: str,  vep_cache: str,  vep_assembly: str, process: int):
    vep_command = '''
            vep --offline --refseq --use_given_ref \
                --dir_cache ''' + vep_cache + ''' \
                --species "homo_sapiens" \
                --assembly ''' + vep_assembly + ''' \
                --fork ''' + str(process) + ''' \
                --hgvs --hgvsg --canonical --symbol \
                --distance 0 \
                --exclude_predicted \
                --flag_pick \
                --lookup_ref \
                --force \
                --input_file ''' + vep_input + ''' \
                --output_file ''' + vep_output + ''' --no_stats \
                --numbers \
                --tab --fields "Uploaded_variation,SYMBOL,Feature,CANONICAL,PICK,EXON,INTRON,Consequence,CDS_position"
        '''
    if os.system(vep_command) > 0:
        raise Exception('run VEP FAILED')
